import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib的字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体为微软雅黑
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题


# 接受gamma作为参数
def simulate_disease_spread(gamma, N, beta, sigma, delta, alpha, rho, dt, T):
    # 初始化
    S0, V0, E0, I0, R0, D0 = N - 1, 0, 0, 1, 0, 0  # 初始条件
    S, V, E, I, R, D = [S0], [V0], [E0], [I0], [R0], [D0]
    np.random.seed(0)

    # Euler-Maruyama方法迭代更新
    for t in np.arange(dt, T, dt):
        dW_S, dW_V, dW_E, dW_I, dW_R, dW_D = np.random.normal(0, np.sqrt(dt), 6)
        S_new = S[-1] + (-beta * S[-1] * (I[-1] + E[-1]) / N - gamma * S[-1]) * dt + sigma * S[-1] * dW_S
        V_new = V[-1] + (gamma * S[-1] - sigma * V[-1] * (I[-1] + E[-1]) / N) * dt + sigma * V[-1] * dW_V
        E_new = E[-1] + (beta * S[-1] * (I[-1] + E[-1]) / N + sigma * V[-1] * (I[-1] + E[-1]) / N - delta * E[
            -1]) * dt + sigma * E[-1] * dW_E
        I_new = I[-1] + (delta * E[-1] - (alpha + rho) * I[-1]) * dt + sigma * I[-1] * dW_I
        R_new = R[-1] + alpha * I[-1] * dt + sigma * R[-1] * dW_R
        D_new = D[-1] + rho * I[-1] * dt + sigma * D[-1] * dW_D
        S.append(S_new)
        V.append(V_new)
        E.append(E_new)
        I.append(I_new)
        R.append(R_new)
        D.append(D_new)

    return S, V, E, I, R, D


# 主程序
if __name__ == "__main__":
    # 要比较的gamma值列表
    gamma_values = [0.05, 0.1, 0.15]

    # 定义模拟参数
    N = 10000  # 总人口
    beta = 0.2  # 感染率
    sigma = 0.01  # 随机波动率
    delta = 0.1  # 暴露者转化为感染者的速率
    alpha = 0.05  # 感染者恢复率
    rho = 0.01  # 感染者死亡率
    dt = 0.1  # 时间步长
    T = 100  # 总时间

    # 初始化结果列表
    results = []

    # 对每个gamma值
    for gamma in gamma_values:
        S, V, E, I, R, D = simulate_disease_spread(gamma, N, beta, sigma, delta, alpha, rho, dt, T)
        results.append((S, V, E, I, R, D))

        # 可视化分析：干预措施效果对比图
    plt.figure(figsize=(12, 8))

    for i, (S, V, E, I, R, D) in enumerate(results):
        plt.plot(np.arange(0, T, dt), I, label=f'传染者 (gamma={gamma_values[i]})')

    plt.xlabel('时间')
    plt.ylabel('传染者人数')
    plt.title('不同疫苗接种率下的传染者数量对比')
    plt.legend()
    plt.grid(True)
    plt.show()